# 非定常的Stokes方程 新版本 重新设计求解过程
from mpi4py import MPI
import numpy as np

import ufl
from basix.ufl import element, mixed_element
from dolfinx import default_real_type, fem, la
from dolfinx.fem import (
    Constant,
    Function,
    dirichletbc,
    form,
    functionspace,
)
from petsc4py import PETSc
from dolfinx.fem.petsc import LinearProblem
from dolfinx.io import XDMFFile
from dolfinx.mesh import CellType, create_rectangle, locate_entities_boundary
from ufl import div, dx, grad, inner, TestFunction, TrialFunction


# He Xiaoming 7.5.1 Example

def true_velocity(x, t):
    return [u_x(x, t), u_y(x, t)]


def true_pressure(x, t):
    return pressure(x, t)


def u_x(x, t):
    """
    x方向速度的真实解
    """
    return (x[0]**2 * x[1]**2 + np.exp(-x[1])) * np.cos(2 * np.pi * t)


def u_y(x, t):
    """
    y方向速度的真实解
    """
    return (-2/3 * x[0] * x[1]**3 + 2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * t)


def u_init(x):
    """
    返回u: IC
    """
    t1 = (x[0]**2 * x[1]**2 + np.exp(-x[1]))
    t2 = -2/3 * x[0] * x[1]**3 + 2 - np.pi * np.sin(np.pi * x[0])
    return [t1, t2]


def p_init_expr(x):
    """
    返回p: IC
    """
    return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1])


def p_end_expr(x):
    """
    返回p: IC
    """
    return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1]) * np.cos(2 * np.pi * T)


def pressure(x, t):
    """
    压强的真实解
    """
    return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1]) * np.cos(2 * np.pi * t)


# ufl格式
def f_x(x, t):
    """
    f的x分量
    """
    term1 = -2 * ufl.pi * (x[0]**2 * x[1]**2 +
                           ufl.exp(-x[1])) * ufl.sin(2 * ufl.pi * t)
    term2 = (
        -2 * nu * x[0]**2
        - 2 * nu * x[1]**2
        - nu * ufl.exp(-x[1])
        + ufl.pi**2 * ufl.cos(ufl.pi * x[0]) * ufl.cos(2 * ufl.pi * x[1])
    ) * ufl.cos(2 * ufl.pi * t)
    return term1 + term2


# He给出的数值算例不正确 第二项少了cos(2*pi*t)
def f_y(x, t):
    """
    f的y分量
    """
    term1 = -2 * ufl.pi * (-2/3 * x[0] * x[1]**3 + 2 - ufl.pi *
                           ufl.sin(ufl.pi * x[0])) * ufl.sin(2 * ufl.pi * t)
    term2 = (
        4 * nu * x[0] * x[1]
        - nu * ufl.pi**3 * ufl.sin(ufl.pi * x[0])
        + 2 * ufl.pi * (2 - ufl.pi * ufl.sin(ufl.pi *
                        x[0])) * ufl.sin(2 * ufl.pi * x[1])
    ) * ufl.cos(2 * ufl.pi * t)
    return term1 + term2


def f(x, t):
    """
    返回f(list)
    """
    return [f_x(x, t), f_y(x, t)]


def D(u):
    return 0.5 * (ufl.nabla_grad(u) + ufl.nabla_grad(u).T)


def get_norm(t: float):
    """
    传入时间t
    """
    # 计算误差
    u_true = Function(V)  # 实例化空间
    p_true = Function(Q)

    u_true.interpolate(lambda x: np.array([u_x(x, t), u_y(x, t)]))

    def p_current(x):  # 只能这样设置时间 比较丑陋
        return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1]) * np.cos(2 * np.pi * t)
    p_true.interpolate(p_current)

    u_error = Function(Q)
    u_error.x.array[:] = u_n.x.array - u_true.x.array
    p_error = Function(S)
    p_error.x.array[:] = p_n.x.array - p_true.x.array

    # Compute norms of the u and p vectors
    u_comm = u_n.function_space.mesh.comm
    # type = dolfinx.fem.forms.Form
    u_h1_error = form((grad(u_error)**2 + u_error**2) * ufl.dx)
    u_l2_error = form(u_error**2 * ufl.dx)
    u_h1_E = np.sqrt(u_comm.allreduce(
        fem.assemble_scalar(u_h1_error), MPI.SUM))
    u_l2_E = np.sqrt(u_comm.allreduce(
        fem.assemble_scalar(u_l2_error), MPI.SUM))  # 并行求和 返回float

    p_comm = p_n.function_space.mesh.comm
    # type = dolfinx.fem.forms.Form
    p_h1_error = form((grad(p_error)**2 + p_error**2) * ufl.dx)
    p_l2_error = form(p_error**2 * ufl.dx)
    p_h1_E = np.sqrt(p_comm.allreduce(
        fem.assemble_scalar(p_h1_error), MPI.SUM))
    p_l2_E = np.sqrt(p_comm.allreduce(
        fem.assemble_scalar(p_l2_error), MPI.SUM))  # 并行求和 返回float

    # Linf 比较节点值 离散后区域无法计算理论最大值
    u_linf_E, p_linf_E = la.norm(
        u_error.x, la.Norm.linf), la.norm(p_error.x, la.Norm.linf)

    if MPI.COMM_WORLD.rank == 0:
        print(
            f"H1 norm for velocity: {u_h1_E}, pressure: {p_h1_E}.")
        print(
            f"L2 norm for velocity: {u_l2_E}, pressure: {p_l2_E}.")
        print(
            f"Linf norm for velocity: {u_linf_E}, pressure: {p_linf_E}.")


# Params
n_list = [
    [8, 2],
    [16, 4],
    [32, 8],
    [64, 16],
    [128, 32],
]
# idx = int(input("请输入网格划分的索引: "))
idx = 1
nu = 1
t = 0  # 初始时刻
T = 1  # 终止时刻
time_steps = 64  # 时间步
dt = T / time_steps

msh = create_rectangle(
    MPI.COMM_WORLD, [np.array([0, -0.25]), np.array(
        [1, 0])], n=n_list[idx], cell_type=CellType.triangle
)
x = ufl.SpatialCoordinate(msh)
print(f"Mesh = {n_list[idx]}")

# mixed element space
P2 = element(
    "Lagrange", msh.basix_cell(), degree=2, shape=(msh.geometry.dim,), dtype=default_real_type
)
P1 = element("Lagrange", msh.basix_cell(), degree=1, dtype=default_real_type)
mel = mixed_element([P2, P1])  # 定义混合元空间 分别对应速度 压力
ME = functionspace(msh, mel)  # 实例化混合元空间

Q, _ = ME.sub(0).collapse()
S, _ = ME.sub(1).collapse()
u_bc = Function(Q)  # 插值当前时间边界条件
p_bc = Function(S)  



# 自由度
facets = locate_entities_boundary(
    msh,
    dim=(msh.topology.dim - 1),
    marker=lambda x: np.isclose(x[0], 0.0) | np.isclose(
        x[0], 1.0) | np.isclose(x[1], -0.25) | np.isclose(x[1], 0.0),
)
v_dofs = fem.locate_dofs_topological((ME.sub(0), Q), 1, facets)
# 压力梯度是影响流体流动的关键 压力的绝对值并不会影响（物理角度）
def marker(x): return np.isclose(x[0], 0.0) & np.isclose(x[1], -0.25)
p_dofs = fem.locate_dofs_geometrical((ME.sub(1), S), marker)


# 前一步的计算结果
u_n = Function(Q)
p_n = Function(S)
u_n.interpolate(u_init)
p_n.interpolate(p_init_expr)

# 解函数
sol = Function(ME)
sol.x.array[:] = 0.0
u, p = ufl.split(sol)  # trial: u, p; test: v, q
v, q = ufl.TestFunctions(ME)  # Test Functions

# 配置求解器


while t < T:
    u_bc.interpolate(lambda x: true_velocity(x, t))
    p_bc.interpolate(lambda x: true_pressure(x, t))
    bc_pressure = dirichletbc(p_bc, p_dofs, ME.sub(1))
    bc_vel = dirichletbc(u_bc, v_dofs, ME.sub(0))  # 速度边界条件
    bcs = [bc_vel, bc_pressure]
    f_expr = ufl.as_vector([f_x(x, t), f_y(x, t)])
    # 弱形式
    # a = 2 * nu * inner(D(u), D(v)) * dx
    # a += inner(u, v) / dt * dx -inner(p, div(v)) * dx
    # a += inner(div(u), q) * dx
    # L = inner(f_expr + u_n/dt, v) * dx
    F = (inner(u, v)/dt * dx + 2*nu*inner(D(u), D(v))*dx
         - inner(p, div(v))*dx + inner(div(u), q)*dx
         - inner(f_expr + u_n/dt, v)*dx)

    # 分离双线性形式和线性形式
    a = lhs(F)
    L = rhs(F)
    problem = LinearProblem(a, L, bcs=bcs, petsc_options={
                            "ksp_type": "preonly", "pc_type": "lu"})
    sol = problem.solve()
    
    u_n.x.array[:] = sol.sub(0).collapse().x.array
    p_n.x.array[:] = sol.sub(1).collapse().x.array
    u_n.x.scatter_forward()
    p_n.x.scatter_forward()
    get_norm(t)
    t += dt
    